import {
  cloneElement,
  forwardRef,
  memo,
  useLayoutEffect,
  useRef,
  useState,
} from "react";
import useLatestRef from "@/lib/hooks/useLatestRef";
import useMergeRefs from "@/lib/hooks/useMergeRefs";
import clsm from "@/lib/utils/clsm";
import { createPortal } from "react-dom";

type TransitionClasses = {
  beforeEnterClass: string;
  enteringClass: string;
  enteredClasss: string;
  beforeLeaveClass: string;
  leavingClass: string;
  leftClass: string;
};

type TransitionProps = {
  children: JSX.Element;
  show?: boolean;
  unmountOnHide?: boolean;
  mountOnBody?: boolean;
  appear?: boolean;
  onLeft?: () => void;
} & Partial<TransitionClasses>;

enum STAGE {
  BEFORE_ENTER = "before-enter",
  ENTERING = "entering",
  ENTERED = "entered",
  BEFORE_LEAVE = "before-leave",
  LEAVING = "leaving",
  LEFT = "left",
}

const Transition = forwardRef<any, TransitionProps>(
  function Transition(props, ref) {
    const {
      children,
      show,
      unmountOnHide = false,
      mountOnBody = false,
      appear,
      beforeEnterClass,
      enteringClass,
      enteredClasss,
      beforeLeaveClass,
      leavingClass,
      leftClass,
    } = props;
    const [stage, setStage] = useState(
      show && !appear ? STAGE.ENTERED : STAGE.LEFT,
    );
    const [unmounted, setUnmounted] = useState(!show && unmountOnHide);
    const stageRef = useLatestRef(stage);
    const propsRef = useLatestRef(props);
    const elRef = useRef<HTMLElement>(null);
    const transitionCallbackRef = useRef<() => void>();

    function next(callback: () => void) {
      const el = elRef.current;
      if (!el) {
        throw new Error("cannot access element reference on children");
      }
      if (transitionCallbackRef.current) {
        el.removeEventListener("transitionend", transitionCallbackRef.current);
      }
      transitionCallbackRef.current = function () {
        callback();
        transitionCallbackRef.current = undefined;
      };
      el.addEventListener("transitionend", transitionCallbackRef.current, {
        once: true,
      });
    }

    useLayoutEffect(() => {
      const { unmountOnHide, onLeft } = propsRef.current;
      const s = stageRef.current;
      if (show) {
        if (s === STAGE.LEFT) {
          if (unmountOnHide) {
            setUnmounted(false);
          }
          requestAnimationFrame(() => {
            setStage(STAGE.BEFORE_ENTER);
            requestAnimationFrame(() => {
              setStage(STAGE.ENTERING);
              next(() => {
                setStage(STAGE.ENTERED);
              });
            });
          });
        } else if (s === STAGE.LEAVING) {
          setStage(STAGE.ENTERING);
          next(() => {
            setStage(STAGE.ENTERED);
          });
        }
      } else {
        if (s === STAGE.ENTERED) {
          setStage(STAGE.BEFORE_LEAVE);
          requestAnimationFrame(() => {
            setStage(STAGE.LEAVING);
            next(() => {
              onLeft?.();
              setStage(STAGE.LEFT);
              unmountOnHide && setUnmounted(true);
            });
          });
        } else if (s === STAGE.ENTERING) {
          setStage(STAGE.LEAVING);
          next(() => {
            onLeft?.();
            setStage(STAGE.LEFT);
            unmountOnHide && setUnmounted(true);
          });
        }
      }
    }, [show]); // eslint-disable-line

    const _ref = useMergeRefs(elRef, ref);

    if (unmounted) {
      return null;
    }

    let className = children.props.className as string;

    if (stage === STAGE.BEFORE_ENTER) {
      className = clsm(className, beforeEnterClass);
    } else if (stage === STAGE.ENTERING) {
      className = clsm(className, enteringClass, enteredClasss);
    } else if (stage === STAGE.ENTERED) {
      className = clsm(className, enteredClasss);
    } else if (stage === STAGE.BEFORE_LEAVE) {
      className = clsm(className, beforeLeaveClass);
    } else if (stage === STAGE.LEAVING) {
      className = clsm(className, leavingClass, leftClass);
    } else {
      className = clsm(className, leftClass);
    }

    const element = cloneElement(children, {
      ...children.props,
      className,
      ref: _ref,
    });

    if (mountOnBody) {
      return createPortal(element, document.body);
    }

    return element;
  },
);

export default memo(Transition);
